///
/// Copyright (C) 2016, Dependable Systems Laboratory, EPFL
/// Copyright (C) 2015-2016, Cyberhaven
/// All rights reserved.
///
/// Licensed under the Cyberhaven Research License Agreement.
///

#include <klee/UserSearcher.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/Support/Path.h>
#include <s2e/ConfigFile.h>
#include <s2e/S2E.h>
#include <s2e/S2EExecutor.h>

#include <random>
#include <sstream>
#include <system_error>

#include <boost/regex.hpp>

#include "RecordSearcher.h"

namespace s2e {
namespace plugins {
namespace records {

S2E_DEFINE_PLUGIN(RecordSearcher, "Record searcher", "", "MultiSearcher", "CUPASearcher");

namespace {

class RecordSearcherState : public PluginState {
public:

    RecordSearcherState() {};

    virtual RecordSearcherState *clone() const {
        RecordSearcherState *ret = new RecordSearcherState(*this);
        return ret;
    }

    static PluginState *factory(Plugin *p, S2EExecutionState *s) {
        return new RecordSearcherState();
    }

    virtual ~RecordSearcherState() {
    }
};
}

RecordSearcher::~RecordSearcher(){
	if(m_libwrapper)
		delete m_libwrapper;
}
void RecordSearcher::initialize() {
	m_libwrapper = new Libwrapper();
	try {
		m_libwrapper->init();
	} catch (...) {
	}
	m_statefilename = "tmp.data";
	m_hasRecord = false;

    ConfigFile *cfg = s2e()->getConfig();

    bool ok = false;

    CorePlugin *plg = s2e()->getCorePlugin();

    plg->onStateFork.connect(sigc::mem_fun(*this, &RecordSearcher::onStateFork));
    plg->onStateKill.connect(sigc::mem_fun(*this, &RecordSearcher::onStateKill));

    // \todo reuse directories from HostFiles.
    m_recordDirectory = cfg->getString(getConfigKey() + ".recordDirectory", "", &ok);
    if (!ok) {
        getWarningsStream() << "Please specify recordDirectory\n";
        exit(-1);
    }

    if (!llvm::sys::fs::exists(m_recordDirectory)) {
        getWarningsStream() << m_recordDirectory << " does not exist\n";
        exit(-1);
    }

    m_cupa = s2e()->getPlugin<CUPASearcher>();
    m_multiSearcher = s2e()->getPlugin<MultiSearcher>();

    m_multiSearcher->registerSearcher("RecordSearcher", this);

    m_cachedState = nullptr;
    m_initialState = nullptr;

    m_backupRecords = cfg->getBool(getConfigKey() + ".backupRecords", true);
    m_enableRecords = cfg->getBool(getConfigKey() + ".enableRecords");
    m_maxRecordStates = cfg->getInt(getConfigKey() + ".maxRecordStates");
    m_parallelRecords = cfg->getBool(getConfigKey() + ".enableParallelRecording", true);
    m_isFirstForkDone = cfg->getBool(getConfigKey() + ".isFirstForkDone", false);
    //m_poolSize = cfg->getInt(getConfigKey() + ".poolSize", 0);
    m_poolSize = 0;
    m_nodeID = cfg->getInt(getConfigKey() + ".nodeID", 0);

    if (m_parallelRecords) {
        plg->onStatesSplit.connect(sigc::mem_fun(*this, &RecordSearcher::onStateSplit));
        plg->onProcessForkComplete.connect(sigc::mem_fun(*this, &RecordSearcher::onProcessForkComplete));
    }

    if (m_backupRecords) {
        m_recordBackupDirectory = s2e()->getOutputFilename("records-backup");
        auto error = llvm::sys::fs::create_directory(m_recordBackupDirectory);
        if (error) {
            getWarningsStream() << "Could not create " << m_recordBackupDirectory << '\n';
            exit(-1);
        }
    }


    std::string m_shareInfoLocationworking = m_recordDirectory
    				+ "/working";
    std::error_code mkdirError =  llvm::sys::fs::create_directories(m_shareInfoLocationworking);
	if (mkdirError) {
		std::cerr << "Could not create directory " << m_shareInfoLocationworking
				<< " error: " << mkdirError << '\n';
	}
	std::string m_shareInfoLocationfinished = m_recordDirectory
			+ "/finished";
	mkdirError =  llvm::sys::fs::create_directories(m_shareInfoLocationfinished);
	if (mkdirError) {
		std::cerr << "Could not create directory " << m_shareInfoLocationfinished
				<< " error: " << mkdirError << '\n';
	}
}

void RecordSearcher::switchToCUPA() {
    m_multiSearcher->selectSearcher("CUPASearcher");
}

void RecordSearcher::switchToRecordSearcher() {
    m_multiSearcher->selectSearcher("RecordSearcher");
}

/// We want to keep state 0 in all S2E instances in
/// order to be able to fetch records in parallel.
void RecordSearcher::onStateSplit(klee::StateSet &parent, klee::StateSet &child) {
    if (m_initialState) {
        parent.insert(m_initialState);
        child.insert(m_initialState);
    }
}

void RecordSearcher::onProcessForkComplete(bool isChild) {
    if (isChild) {

    }
}

void RecordSearcher::onStateFork(S2EExecutionState *oldState, const std::vector<S2EExecutionState *> &newStates,
                               const std::vector<klee::ref<klee::Expr>> &) {
    // Initial state becomes record state when it forks after getting record file
    if (oldState == m_initialState) {
    	m_initialState->symbolics.clear();
    	m_initialState->m_forkrecord.clear();
    	m_initialState->m_forkPoints.clear();

        for (auto it : newStates) {
            if (it != oldState) {
            	if( m_hasRecord ){
            		m_hasRecord = false;
            		// make state as replaying
            		std::stringstream ss;
					ss << m_recordDirectory << "/working/" << m_statefilename << ".working";
					std::ifstream is(ss.str().c_str(), std::ifstream::binary);
					transferstate ts;
					getLibwrapper()->deserialize(ts, is);
					is.close();
					it->symbolics.clear();
					it->m_forkrecord.clear();
					it->m_forkPoints.clear();
					it->m_forkrecord4repaly = ts.m_forkrecord;
					it->m_concreteAddressSequence.clear();
					it->m_concreteAddressSequence4repaly= ts.m_concreteAddressSequence;
					if (it->m_forkrecord4repaly.size() > 0) {
						it->m_replaying = true;
					}
					it->m_rgdMode = true;
					it->m_statefilename = m_statefilename;
            		m_recordStates.insert(it);
            	}else{
            		it->symbolics.clear();
            		it->m_forkrecord.clear();
            		it->m_forkPoints.clear();
            		m_recordStates.insert(it);
            	}
            }
        }
    }
}

void RecordSearcher::onStateKill(S2EExecutionState *state) {
	if(state->m_should_del || state->m_replaying){
		return;
	}
	if (state->m_statefilename.length() > 1) {
		std::stringstream ssocu;
		ssocu << m_recordDirectory << "/working/" << state->m_statefilename
				<< ".working";
		std::stringstream ssfinish;
		ssfinish << m_recordDirectory << "/finished/" << m_nodeID << "-"
				<< state->m_statefilename;
		try {
			std::error_code error;
			llvm::raw_fd_ostream *out = new llvm::raw_fd_ostream(
					ssfinish.str().c_str(), error,
					llvm::sys::fs::F_None);

			if (!out || error) {
				throw;
			}
			std::ifstream in(ssocu.str().c_str());
			char c;
			while (in.get(c)) {
				(*out) << c;
			}
			out->flush();
			out->close();
			delete out;
			remove(ssocu.str().c_str());
		} catch (...) {
			getWarningsStream(state) << "Could not write replay state finished file: " << ssfinish.str().c_str() << "\n";
		}
	} else {
		std::stringstream ssfinish;
		ssfinish << m_recordDirectory << "/finished/" << m_nodeID
				<< "-" << state->getID();
		try {
			std::ofstream out(ssfinish.str().c_str(),
					std::ofstream::binary);
			transferstate ts;
			ts.m_forkrecord = state->m_forkrecord;
			ts.m_forkrecord4repaly = state->m_forkrecord4repaly;
			ts.m_forkPoints = state->m_forkPoints;
			ts.m_concreteAddressSequence = state->m_concreteAddressSequence;
			ts.m_concreteAddressSequence4repaly = state->m_concreteAddressSequence4repaly;
			getLibwrapper()->serialize(ts, out);
			out.flush();
			out.close();
		} catch (...) {
			getWarningsStream(state) << "Could not write nomal state finished file: " << ssfinish.str().c_str() << "\n";
		}
	}
}

void RecordSearcher::update(klee::ExecutionState *current, const klee::StateSet &addedStates,
                          const klee::StateSet &removedStates) {
    S2EExecutionState *cs = dynamic_cast<S2EExecutionState *>(current);
    if (!m_initialState) {
        if (cs->getID() == 0) {
            m_initialState = cs;
        }
    }
    std::set<klee::ExecutionState*> terminateStates;
    for (auto addedState : addedStates) {
        S2EExecutionState *es = dynamic_cast<S2EExecutionState *>(addedState);
        if(es->m_should_del){
        	terminateStates.insert(addedState);
        	continue;
        }
        if((m_states.size()-2 > m_poolSize && m_initialState != es && cs->m_should_del) || (m_states.size()-1 > m_poolSize && m_initialState != es && !cs->m_should_del)){ //
        	// serialize
        	std::stringstream ss;
			ss << m_recordDirectory << "/" << m_nodeID << "-"
					<< es->getID();
			std::ofstream os(ss.str().c_str(), std::ofstream::binary);

			transferstate ts;
			ts.m_forkrecord = es->m_forkrecord;
			ts.m_forkrecord4repaly = es->m_forkrecord4repaly;
			ts.m_forkPoints = es->m_forkPoints;
			ts.m_concreteAddressSequence = es->m_concreteAddressSequence;
			ts.m_concreteAddressSequence4repaly = es->m_concreteAddressSequence4repaly;
			getLibwrapper()->serialize(ts, os);
			os.flush();
			os.close();

			es->m_should_del = true;
			terminateStates.insert(addedState);
        }else{
        	m_states.insert(es);
        }
    }

    for (auto removedState : removedStates) {
        S2EExecutionState *es = dynamic_cast<S2EExecutionState *>(removedState);

        // This can only happen if state 0 dies for some reason
        if (es == m_initialState) {
            s2e_warn_assert(cs, false, "Initial state no longer exists, record look up is not possible");
            m_initialState = nullptr;
        }

        if (es == m_cachedState) {
            m_cachedState = nullptr;
        }

        m_recordStates.erase(es);
        m_states.erase(es);
    }

    for (auto terminateState : terminateStates) {
    	S2EExecutionState *es = dynamic_cast<S2EExecutionState *>(terminateState);
        if (es == m_cachedState) {
            m_cachedState = nullptr;
        }
    	m_recordStates.erase(es);
        m_states.erase(es);
        if(!es->isZombie()){
        	s2e()->getExecutor()->terminateState(*terminateState,"terminate writen-out state.");
        }
    }

    if(cs && cs->m_should_del == true && !cs->isZombie()){
        if (cs == m_cachedState) {
            m_cachedState = nullptr;
        }
    	m_recordStates.erase(cs);
        m_states.erase(cs);

		s2e()->getExecutor()->terminateState(*cs,"terminate replay state.");
	}
}

klee::ExecutionState &RecordSearcher::selectState() {
    if (m_cachedState) {
         return *m_cachedState;
     }

    auto it = m_recordStates.begin();
    if (it != m_recordStates.end()) {
        m_cachedState = *it;
        return *m_cachedState;
    }

    if (m_states.size()-1 < m_poolSize && (m_initialState && m_hasRecord)) {
        return *m_initialState;
    }

    auto itx = m_states.begin();
    if(*itx == m_initialState){
    	itx++;
    	if (itx != m_states.end()) {
			m_cachedState = *itx;
			return **itx;
		} else {
			return *m_initialState;
		}
    }else{
		m_cachedState = *itx;
		return **itx;
    }

}

bool RecordSearcher::empty() {
    return m_states.empty();
}

void RecordSearcher::handleGetRecordFile(S2EExecutionState *state, S2E_RECORDSEARCHER_COMMAND &cmd) {

	bool should_fork = false;
	if (!m_hasRecord) {
		int maxfile = 0;
		file_dlib task;
		try {
			std::vector<file_dlib> tasks = getLibwrapper()->getfiles(
					m_recordDirectory);
			maxfile = tasks.size();
			int selectindex = maxfile - 1;
			if (maxfile > 0) {
				selectindex = rand() % maxfile;
				task = tasks[selectindex];
				m_statefilename = task.name;
			} else {
			}
		} catch (...) {
		}

		if (maxfile > 0) {
			try {
				std::stringstream ss;
				ss << m_recordDirectory << "/" << m_statefilename;
				std::stringstream ssocu;
				ssocu << m_recordDirectory << "/working/" << m_statefilename
						<< ".working";
				std::error_code error;
				llvm::raw_fd_ostream *out = new llvm::raw_fd_ostream(
						ssocu.str().c_str(), error,
						llvm::sys::fs::F_None);
				if (!out|| error) {
					throw;
				}
				std::ifstream in(ss.str().c_str());
				in.seekg(std::ios::beg);
				char c;
				while (in.get(c)) {
					(*out) << c;
				}
				out->flush();
				out->close();
				in.close();
				delete out;

				//delete record
				try {
					remove(ss.str().c_str());
				} catch (...) {
				}
				should_fork = true;
			} catch (...) {
			}
		}
	}

	if((!m_hasRecord && should_fork)||!m_isFirstForkDone){
		m_isFirstForkDone = true;
		cmd.GetFile.Result = 2;
	    unsigned length = std::min(cmd.GetFile.FileNameSizeInBytes, m_statefilename.length() + 1);
	    if (!state->mem()->write(cmd.GetFile.FileName, m_statefilename.c_str(), length)) {
	        getDebugStream(state) << "Could not write " << m_statefilename << " to " << hexval(cmd.GetFile.FileName)
	                              << "\n";
	        exit(-1);
	    }
	    if(!m_hasRecord && should_fork){
	    	m_hasRecord = true;
	    }
	    return;
	}else{
		cmd.GetFile.Result = 0;
		return;
	}

}

void RecordSearcher::enableRecords(bool enable) {
    if (enable && !m_enableRecords) {
        getDebugStream() << "Enabling records\n";
    } else if (!enable && m_enableRecords) {
        getDebugStream() << "Disabling records\n";
    }

    m_enableRecords = enable;
}


void RecordSearcher::handleOpcodeInvocation(S2EExecutionState *state, uint64_t guestDataPtr, uint64_t guestDataSize) {
    S2E_RECORDSEARCHER_COMMAND command;

    if (guestDataSize != sizeof(command)) {
        getWarningsStream(state) << "S2E_RECORDSEARCHER_COMMAND: mismatched command structure size " << guestDataSize
                                 << "\n";
        exit(-1);
    }

    if (!state->mem()->read(guestDataPtr, &command, guestDataSize)) {
        getWarningsStream(state) << "S2E_RECORDSEARCHER_COMMAND: could not read transmitted data\n";
        exit(-1);
    }

    switch (command.Command) {
        case RECORD_GET_RECORD_FILE: {
            handleGetRecordFile(state, command);
            if (!state->mem()->write(guestDataPtr, &command, sizeof(command))) {
                getWarningsStream(state) << "Could not write to guest memory\n";
            }
        } break;

        case RECORD_ENABLE_SEARCHER: {
            switchToRecordSearcher();
        } break;

        case RECORD_DONE: {
            switchToCUPA();
        } break;

        default: {
            getWarningsStream(state) << "Invalid command " << hexval(command.Command) << "\n";
            exit(-1);
        }
    }
}

} // namespace records
} // namespace plugins
} // namespace s2e
